package common

import (
	"context"
	"github.com/coreos/etcd/clientv3"
)

type JobLock struct {
	kv       clientv3.KV
	lease    clientv3.Lease
	leaseId  clientv3.LeaseID
	jobName  string
	cancel   context.CancelFunc
	isLocked bool
}

func InitJobLock(jobName string, kv clientv3.KV, lease clientv3.Lease) (jobLock *JobLock) {

	jobLock = &JobLock{
		kv:       kv,
		lease:    lease,
		jobName:  jobName,
		isLocked: false,
	}
	return
}

func (jk *JobLock) TryLock() (err error) {
	var (
		leaseRes *clientv3.LeaseGrantResponse
		//resCh    <-chan *clientv3.LeaseKeepAliveResponse
		txn    clientv3.Txn
		key    string
		txnRes *clientv3.TxnResponse
	)
	if leaseRes, err = jk.lease.Grant(context.TODO(), 5); err != nil {
		return
	}

	ctx, cancel := context.WithCancel(context.TODO()) //用户自动去掉续租

	if _, err = jk.lease.KeepAlive(ctx, leaseRes.ID); err != nil {
		goto FAILED
	}
	txn = jk.kv.Txn(context.TODO())
	key = ETCD_JOBS_LOCK + jk.jobName
	txn.If(clientv3.Compare(clientv3.CreateRevision(key), "=", 0)).
		Then(clientv3.OpPut(key, "", clientv3.WithLease(leaseRes.ID))).Else(clientv3.OpGet(key))
	if txnRes, err = txn.Commit(); err != nil {
		err = ERR_LOCKED
		goto FAILED
	}
	if !txnRes.Succeeded {
		err = ERR_LOCKED
		goto FAILED
	}
	jk.isLocked = true
	jk.cancel = cancel
	jk.leaseId = leaseRes.ID
	return

FAILED:

	cancel()
	jk.lease.Revoke(context.TODO(), leaseRes.ID)
	return
}

func (jk *JobLock) Unlock() {
	if jk.isLocked {
		jk.cancel()
		jk.lease.Revoke(context.TODO(), jk.leaseId)
	}
}
